import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
import pickle
import logging  # 用于记录日志

# 加载图像
image = cv2.imread('Sandstone_1.tif')
image1 = cv2.imread('Sandstone_1_segment.tif')
print('图像形状:', image.shape)
image1_gray = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY) # 将图像由BGR变为灰度图
features = image.reshape(-1, 3)  # 将图像数据重塑为一个二维数组，以便于存储X和y
labels = image1_gray.flatten()  # 将灰度图像数据展平为一维数组，每个元素表示一个像素的灰度值
print('完成从砂岩截面图1及其对应分区中获取X和y')
X_train, X_test, y_train, y_test = train_test_split(features, labels, test_size=0.3, random_state=50)  # 划分训练集和测试集
clf = RandomForestClassifier()  # 创造随机森林模型
clf.fit(X_train, y_train)
predictions = clf.predict(X_test)  # 在测试集上进行预测
print('完成train_test_split')
print('完成随机森林模型clf的训练')
accuracy = np.mean(predictions == y_test)  # 计算预测结果与真实标签之间的准确率
print('准确率:', accuracy)
# 保存模型到硬盘
with open('clf_model.pkl', 'wb') as file:
    pickle.dump(clf, file)
print('已保存保存随机森林模型clf到硬盘')